#include "undistortion.h"

Undistortion::Undistortion(const std::string& path) {
	std::string filename = path +"calibration/left.yaml";
   cv::FileStorage fs(filename, cv::FileStorage::READ);
   if (!fs.isOpened()) return;
        
   int width, height;
   fs["image_width"] >> width;
   fs["image_height"] >> height;
   image_size = cv::Size(width, height);
        
   fs["camera_matrix"] >> K;
   fs["distortion_coefficients"] >> D;
   fs["rectification_matrix"] >> R;
   fs["projection_matrix"] >> P;

   
   fs.release();

   std::cout << " K : \n" << K << std::endl;     

   image_width_scale_ratio_ = image_size.width / FRAME_WIDTH;
   image_height_scale_ratio_ = image_size.height / FRAME_HEIGH;
}

void Undistortion::ImageResize(cv::Mat& image, const cv::Mat& frame) {
	// cv::resize(frame, image, cv::Size(FRAME_WIDTH, FRAME_HEIGH));
	image = frame;
  	if(image.channels() != 1) {
    	cv::cvtColor(image, image, cv::COLOR_RGB2GRAY);
  	}
}

void Undistortion::UnDistortionKeyPoint(cv::KeyPoint& un_kp, cv::KeyPoint& resize_kp, const cv::KeyPoint& kp) {
	 resize_kp = kp;
	 if (image_height_scale_ratio_ != 1.0 || image_width_scale_ratio_ != 1.0) {
		resize_kp.pt.x = kp.pt.x * image_width_scale_ratio_;
		resize_kp.pt.y = kp.pt.y * image_height_scale_ratio_;
	 }else {
		resize_kp.pt = kp.pt;
	 }

	 // 检查输入有效性
    if (K.empty() || D.empty()) {
        std::cerr << "Error: Camera matrix or distortion coefficients are empty!" << std::endl;
        return;
    }
    
    std::vector<cv::Point2f> distorted_points = {resize_kp.pt};
    std::vector<cv::Point2f> undistorted_points;
    
    try {
        cv::undistortPoints(distorted_points, undistorted_points, K, D);

        double fx = K.at<double>(0,0);
        double fy = K.at<double>(1,1);
        double cx = K.at<double>(0,2);
        double cy = K.at<double>(1,2);
            
        undistorted_points[0].x = undistorted_points[0].x * fx + cx;
        undistorted_points[0].y = undistorted_points[0].y * fy + cy;
       
        // 检查去畸变后的点是否在图像范围内（如果提供了图像尺寸）
        if (image_size.width > 0 && image_size.height > 0) {
            if (undistorted_points[0].x < 0 || undistorted_points[0].x >= image_size.width ||
                undistorted_points[0].y < 0 || undistorted_points[0].y >= image_size.height) {
                std::cout << "Warning: Undistorted point is outside image boundaries!" << std::endl;
            }
        }
        
        un_kp = resize_kp;  // 复制所有属性
        un_kp.pt = undistorted_points[0];
        
        return;
        
    } catch (const cv::Exception& e) {
        std::cerr << "OpenCV exception in UnDistortionKeyPoint: " << e.what() << std::endl;
        return;
    }

}

void Undistortion::UnDistortionKeyPoint(cv::KeyPoint& un_kp, const cv::KeyPoint& kp) {
    // 检查输入有效性
    if (K.empty() || D.empty()) {
        std::cerr << "Error: Camera matrix or distortion coefficients are empty!" << std::endl;
        return;
    }
    
    std::vector<cv::Point2f> distorted_points = {kp.pt};
    std::vector<cv::Point2f> undistorted_points;
    
    try {
        cv::undistortPoints(distorted_points, undistorted_points, K, D);

        double fx = K.at<double>(0,0);
        double fy = K.at<double>(1,1);
        double cx = K.at<double>(0,2);
        double cy = K.at<double>(1,2);
            
        undistorted_points[0].x = undistorted_points[0].x * fx + cx;
        undistorted_points[0].y = undistorted_points[0].y * fy + cy;
       
        // 检查去畸变后的点是否在图像范围内（如果提供了图像尺寸）
        if (image_size.width > 0 && image_size.height > 0) {
            if (undistorted_points[0].x < 0 || undistorted_points[0].x >= image_size.width ||
                undistorted_points[0].y < 0 || undistorted_points[0].y >= image_size.height) {
                std::cout << "Warning: Undistorted point is outside image boundaries!" << std::endl;
            }
        }
        
        un_kp = kp;  // 复制所有属性
        un_kp.pt = undistorted_points[0];
        
        return;
        
    } catch (const cv::Exception& e) {
        std::cerr << "OpenCV exception in UnDistortionKeyPoint: " << e.what() << std::endl;
        return;
    }
}

Eigen::Vector3d Undistortion::NormalizedPoint(const cv::KeyPoint& undistorted_kp) {
    double fx = K.at<double>(0,0);
    double fy = K.at<double>(1,1);
    double cx = K.at<double>(0,2);
    double cy = K.at<double>(1,2);

    Eigen::Vector3d normalized_point;
    normalized_point << (undistorted_kp.pt.x-cx)/fx, (undistorted_kp.pt.y-cy)/fy, 1.0;

    return normalized_point;
}

void Undistortion::DistortionKeyPoint(cv::KeyPoint& kp,const cv::KeyPoint& un_kp) {
    if (K.empty() || D.empty()) {
        std::cerr << "Error: Camera matrix or distortion coefficients are empty!" << std::endl;
        return;
    }
    
    // 将归一化坐标系的点转换到像素坐标系
    double fx = K.at<double>(0,0);
    double fy = K.at<double>(1,1);
    double cx = K.at<double>(0,2);
    double cy = K.at<double>(1,2);
    
    // 将像素坐标转换到归一化坐标系
    std::vector<cv::Point2f> normalized_points;
    cv::Point2f normalized_pt;
    normalized_pt.x = (un_kp.pt.x - cx) / fx;
    normalized_pt.y = (un_kp.pt.y - cy) / fy;
    normalized_points.push_back(normalized_pt);
    
    std::vector<cv::Point2f> distorted_points;
    
    try {
        // 使用 projectPoints 来添加畸变
        std::vector<cv::Point3f> object_points;
        object_points.push_back(cv::Point3f(normalized_pt.x, normalized_pt.y, 1.0f));
        
        cv::Mat rvec = cv::Mat::zeros(3, 1, CV_64FC1);  // 零旋转向量
        cv::Mat tvec = cv::Mat::zeros(3, 1, CV_64FC1);  // 零平移向量
        
        cv::projectPoints(object_points, rvec, tvec, K, D, distorted_points);
        
        // 检查畸变后的点是否在图像范围内
        if (image_size.width > 0 && image_size.height > 0) {
            if (distorted_points[0].x < 0 || distorted_points[0].x >= image_size.width ||
                distorted_points[0].y < 0 || distorted_points[0].y >= image_size.height) {
                std::cout << "Warning: Distorted point is outside image boundaries!" << std::endl;
            }
        }
        
        kp = un_kp;  // 复制所有属性
        kp.pt = distorted_points[0];
        
        return;
        
    } catch (const cv::Exception& e) {
        std::cerr << "OpenCV exception in DistortionKeyPoint: " << e.what() << std::endl;
        return;
    }
}

bool Undistortion::IsInMask(const cv::KeyPoint& kp) {
    return true;
}

Eigen::Matrix3d Undistortion::read_intrinsic() {
    // 校验输入矩阵的有效性（必须是 3x3 单通道矩阵）
    if (K.rows != 3 || K.cols != 3 || K.channels() != 1) {
        throw std::invalid_argument("Input cv::Mat must be a 3x3 single-channel matrix.");
    }

    Eigen::Matrix3d eigen_mat;

    // 根据 cv::Mat 的数据类型提取元素（支持 CV_64F 和 CV_32F）
    if (K.type() == CV_64F) {
        // 双精度（double）矩阵
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                eigen_mat(i, j) = K.at<double>(i, j);  // (行, 列)
            }
        }
    } else if (K.type() == CV_32F) {
        // 单精度（float）矩阵，转换为 double 存储
        for (int i = 0; i < 3; ++i) {
            for (int j = 0; j < 3; ++j) {
                eigen_mat(i, j) = static_cast<double>(K.at<float>(i, j));
            }
        }
    } else {
        throw std::invalid_argument("Input cv::Mat must be of type CV_64F or CV_32F.");
    }

    return eigen_mat;
}